from __future__ import annotations

import ast
import logging
from ast import literal_eval
from enum import Enum
from typing import TYPE_CHECKING, Literal
from uuid import uuid4

import pandas as pd
from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode
from pybox.base import BasePyBoxManager  # noqa: TCH002

from tablegpt.agent.file_reading.data_normalizer import (
    get_data_normalize_chain,
    get_table_reformat_chain,
    wrap_normalize_code,
)
from tablegpt.errors import NoAttachmentsError
from tablegpt.tools import IPythonTool, markdown_console_template
from tablegpt.translation import create_translator

if TYPE_CHECKING:
    from pathlib import Path

    from langchain_core.language_models import BaseLanguageModel

logger = logging.getLogger(__name__)


class Stage(Enum):
    UPLOADED = 0
    INFO_READ = 1
    HEAD_READ = 2


class AgentState(MessagesState):
    # The message that we received from the user, act as an entry point
    entry_message: BaseMessage
    processing_stage: Stage
    # This is a bit of a hack to pass parent id to the agent state
    # But it act as the group id of all messages generated by the agent
    parent_id: str | None


ENCODER_INPUT_SEG_NUM = 2


def create_file_reading_workflow(
    llm: BaseLanguageModel,
    pybox_manager: BasePyBoxManager,
    *,
    workdir: Path | None = None,
    session_id: str | None = None,
    nlines: int | None = None,
    normalize_llm: BaseLanguageModel | None = None,
    locale: str | None = None,
    verbose: bool = False,
):
    """Create a workflow for reading and processing files using an agent-based approach.

    Args:
        llm (Runnable): The primary language model for processing user input.
        pybox_manager (BasePyBoxManager): A Python code sandbox delegator.
        workdir (Path | None, optional): The working directory for `pybox` operations. Defaults to None.
        session_id (str | None, optional): An optional session identifier used to associate with `pybox`. Defaults to None.
        nlines (int | None, optional): The number of lines to display from the dataset head. Defaults to 5 if not provided.
        normalize_llm (BaseLanguageModel | None, optional): An optional language model used for data normalization. Defaults to None.
        locate (str | None, optional): The locale of the user. Defaults to None.
        verbose (bool, optional): Flag to enable verbose logging for debugging. Defaults to False.

    Returns:
        StateGraph: The compiled state graph representing the workflow for reading and processing files.
    """
    if nlines is None:
        nlines = 5

    # Read the data header into different formats according to the model type.
    model_type = None
    if llm.metadata is not None:
        model_type = llm.metadata.get("model_type")

    translation_chain = None
    if locale is not None:
        translation_chain = create_translator(llm=llm)
    ipython_tool = IPythonTool(pybox_manager=pybox_manager, cwd=workdir, session_id=session_id)
    tool_executor = ToolNode([ipython_tool])

    async def agent_node(state: AgentState) -> dict:
        if state.get("processing_stage", Stage.UPLOADED) == Stage.UPLOADED:
            return await get_df_info(state)
        if state.get("processing_stage", Stage.UPLOADED) == Stage.INFO_READ:
            return get_df_head(state)

        return get_final_answer(state)

    async def generate_normalization_code(state: AgentState) -> str:
        if attachments := state["entry_message"].additional_kwargs.get("attachments"):
            # TODO: we only support one file for now
            filename = attachments[0]["filename"]
        else:
            raise NoAttachmentsError

        var_name = state["entry_message"].additional_kwargs.get("var_name", "df")

        # TODO: refactor the data normalization to langgraph
        content = await ipython_tool.ainvoke(
            input=RAW_TABLE_INFO_CODE.format_map({"filepath": filename, "var_name": f"{var_name}_5rows"})
        )
        raw_table_info = ast.literal_eval(next(x["text"] for x in content if x["type"] == "text"))
        table_reformat_chain = get_table_reformat_chain(llm=normalize_llm)
        reformatted_table = await table_reformat_chain.ainvoke(input={"table": raw_table_info})

        # TODO: Replace pandas dependency with a lightweight alternative or custom implementation.
        if pd.DataFrame(reformatted_table).astype(str).equals(pd.DataFrame(raw_table_info)):
            return ""

        normalize_chain = get_data_normalize_chain(llm=normalize_llm)
        normalization_code: str = await normalize_chain.ainvoke(
            input={
                "table": raw_table_info,
                "reformatted_table": reformatted_table,
            }
        )
        # Add try-except block to catch any errors that may occur during normalization.
        return wrap_normalize_code(var_name, normalization_code)

    async def get_df_info(state: AgentState) -> dict:
        if attachments := state["entry_message"].additional_kwargs.get("attachments"):
            # TODO: we only support one file for now
            filename = attachments[0]["filename"]
        else:
            raise NoAttachmentsError

        var_name = state["entry_message"].additional_kwargs.get("var_name", "df")

        thought = f"我已经收到您的数据文件，我需要查看文件内容以对数据集有一个初步的了解。首先我会读取数据到 `{var_name}` 变量中，并通过 `{var_name}.info` 查看 NaN 情况和数据类型。"  # noqa: RUF001
        if translation_chain is not None:
            thought = await translation_chain.ainvoke(input={"locale": locale, "input": thought})

        read_df_code = f"""# Load the data into a DataFrame
{var_name} = read_df('{filename}')"""

        normalization_code = ""
        if normalize_llm is not None:
            try:
                normalization_code = await generate_normalization_code(state)
            except Exception as e:  # noqa: BLE001
                logger.warning("Failed to generate normalization code: %s", str(e))

        tool_input = f"""{read_df_code}
{normalization_code}
# Remove leading and trailing whitespaces in column names
{var_name}.columns = {var_name}.columns.str.strip()

# Remove rows and columns that contain only empty values
{var_name} = {var_name}.dropna(how='all').dropna(axis=1, how='all')

# Get the basic information of the dataset
{var_name}.info(memory_usage=False)"""

        content = f"{thought}\n```python\n{tool_input}\n```"

        return {
            "messages": [
                AIMessage(
                    id=str(uuid4()),
                    content=content,
                    tool_calls=[
                        {
                            "name": "python",
                            "args": {"query": tool_input},
                            "id": str(uuid4()),
                        }
                    ],
                    additional_kwargs={
                        "parent_id": state["parent_id"],
                        "thought": thought,
                        "action": {
                            "tool": "python",
                            "tool_input": tool_input,
                        },
                        "model_type": model_type,
                    },
                )
            ],
            "processing_stage": Stage.INFO_READ,
        }

    def get_df_head(state: AgentState) -> dict:
        var_name = state["entry_message"].additional_kwargs.get("var_name", "df")

        thought = f"""接下来我将用 `{var_name}.head({nlines})` 来查看数据集的前 {nlines} 行。"""
        if translation_chain is not None:
            thought = translation_chain.invoke(input={"locale": locale, "input": thought})

        # The input visible to the LLM can prevent it from blindly imitating the actions of our encoder.
        default_tool_input = f"""# Show the first {nlines} rows to understand the structure
{var_name}.head({nlines})"""

        # Use the flush parameter to force a refresh of the buffer and return it to multiple text parts
        if model_type == "mm-tabular/markup":
            tool_input = f"""# Show the first {nlines} rows to understand the structure
print({var_name}.head({nlines}), flush=True)
print({var_name}.head(500).to_markdown(), flush=True)"""

        elif model_type == "mm-tabular/contrastive":
            tool_input = f"""# Show the first {nlines} rows to understand the structure
print({var_name}.head({nlines}), flush=True)
print(str(inspect_df({var_name})), flush=True)"""

        else:
            tool_input = default_tool_input

        # The input visible to the LLM can prevent it from blindly imitating the actions of our encoder.
        content = f"{thought}\n```python\n{default_tool_input}\n```"

        return {
            "messages": [
                AIMessage(
                    id=str(uuid4()),
                    content=content,
                    tool_calls=[
                        {
                            "name": "python",
                            "args": {"query": tool_input},
                            "id": str(uuid4()),
                        }
                    ],
                    additional_kwargs={
                        "parent_id": state["parent_id"],
                        "thought": thought,
                        "action": {
                            "tool": "python",
                            "tool_input": default_tool_input,
                        },
                        "model_type": model_type,
                    },
                )
            ],
            "processing_stage": Stage.HEAD_READ,
        }

    def get_final_answer(state: AgentState) -> dict:
        if attachments := state["entry_message"].additional_kwargs.get("attachments"):
            # TODO: we only support one file for now
            filename = attachments[0]["filename"]
        else:
            raise NoAttachmentsError

        text = f"我已经了解了数据集 {filename} 的基本信息。请问我可以帮您做些什么？"  # noqa: RUF001

        if translation_chain is not None:
            text = translation_chain.invoke(input={"locale": locale, "input": text})

        return {
            "messages": [
                AIMessage(
                    id=str(uuid4()),
                    content=text,
                    additional_kwargs={
                        "parent_id": state["parent_id"],
                    },
                )
            ]
        }

    async def tool_node(state: AgentState) -> dict:
        messages: list[ToolMessage] = await tool_executor.ainvoke(state["messages"])
        for message in messages:
            message.additional_kwargs = message.additional_kwargs | {
                "parent_id": state["parent_id"],
                # Hide the execution results of the file upload tool.
                "display": False,
            }
            # TODO: this is very hard-coded to format encoder input like this.
            if (
                model_type in {"mm-tabular/markup", "mm-tabular/contrastive"}
                and len(message.content) == ENCODER_INPUT_SEG_NUM
            ):
                _df_head, _extra = message.content
                table_content = (
                    [literal_eval(_extra["text"])] if model_type == "mm-tabular/contrastive" else [_extra["text"]]
                )

                message.content = [
                    _df_head,
                    {"type": "table", "tables": table_content},
                ]
                message.additional_kwargs["hackfor"] = "encoder"
            for part in message.content:
                if isinstance(part, dict) and part.get("type") == "text":
                    part["text"] = markdown_console_template.format(res=part["text"])
        return {"messages": messages}

    # I cannot use `END` as the literal hint, as:
    #  > Type arguments for "Literal" must be None, a literal value (int, bool, str, or bytes), or an enum value.
    # As `END` is just an intern string of "__end__" (See `langgraph.constants`), So I use "__end__" here.
    def should_continue(state: AgentState) -> Literal["tool_node", "__end__"]:
        # Must have at least one message when entering this router
        last_message = state["messages"][-1]
        if last_message.tool_calls:
            return "tool_node"
        return END

    workflow = StateGraph(AgentState)

    workflow.add_node(agent_node)
    workflow.add_node(tool_node)

    workflow.add_edge(START, "agent_node")
    workflow.add_edge("tool_node", "agent_node")
    workflow.add_conditional_edges("agent_node", should_continue)

    return workflow.compile(debug=verbose)


RAW_TABLE_INFO_CODE = """import numpy as np
from datetime import datetime

{var_name} = read_df('{filepath}', nrows=5, header=None)

# Replace NaN with None and format datetime cells
{var_name} = {var_name}.where({var_name}.notnull(), None).map(
    lambda cell: (
        cell.strftime('%Y-%m-%d') if isinstance(cell, (pd.Timestamp, datetime)) and pd.notnull(cell) else cell
    )
)
# Convert DataFrame to a list of lists
print({var_name}.replace(np.nan, None).values.tolist())"""
